import re
from harvesttext import HarvestText
from kg_back.settings import BASE_DIR


def make_ht():
    ht = HarvestText()
    with open(BASE_DIR.joinpath('api', 'tools', 'new_words.txt'), 'r') as fp:
        new_words = list(fp.readlines())
    ht.add_new_words(new_words)
    return ht


def find_id(name: str, data_list: list):
    for node in data_list:
        if node['name'] == name:
            return node['id']
    return -1


def generate_graph_tuple(ht: HarvestText, doc: str) -> list:
    graph_tuple = []
    doc = re.sub('\\n', '', doc)
    doc = ht.clean_text(doc)
    try:
        for relation in ht.triple_extraction(doc):
            graph_tuple.append(tuple(relation))
        return graph_tuple
    except Exception as e:
        print(e)
        return e


def make_graph_data(graph_tuple: list, ht: HarvestText) -> dict:
    things = set()
    for graph_node in graph_tuple:
        a, v, b = graph_node
        things.add(a)
        things.add(b)
    
    named_entity = ht.named_entity_recognition(' '.join(things))
    categories = [{'name': key} for key in set(named_entity.values())]
    categories.append({'name': '未知'})
    key2categories = named_entity
    nodes, idx = [], 0
    for graph_node in graph_tuple:
        a, v, b = graph_node
        for node in [a, b]:
            nodes.append({
                'category': key2categories.get(node, '未知'),
                'id': idx, 'name': node
            })
            idx += 1

    links = []
    for graph_node in graph_tuple:
        a, v, b = graph_node
        source = find_id(a, nodes)
        target = find_id(b, nodes)
        if source >=0 and target >= 0:
            links.append({'source': source, 'target': target, 'value': v})
    graph_data = {
        'categories': categories,
        'nodes': nodes, 'links': links
    }
    return graph_data

ht = make_ht()

def get_graph_data(doc: str):
    graph_tuple = generate_graph_tuple(ht, doc)
    return make_graph_data(graph_tuple, ht)
